;;Runge-Kutta-Fehlberg method to solve ode
(define scalar-product
  (lambda (h ls)
    (map (lambda (ele) (* h ele)) ls)))
(define add
  (lambda lss
    (if (null? (car lss))
        '()
        (cons (apply + (map car lss))
              (apply add (map cdr lss))))))
(define norm
  (lambda (y)
    (sqrt (apply + (map (lambda (ele) (expt ele 2)) y)))))
(define rkf
  (lambda (function init-time init-val init-step tol end-condition min-step)
    (define f function)
    (define next-state ; a state = (vector t, (y0,y1,..), h)
      (lambda (now-state)
        (let* ([t (vector-ref now-state 0)]
               [y (vector-ref now-state 1)]
               [h (vector-ref now-state 2)]
               [k1 (scalar-product h (f t y))]
               [k2 (scalar-product h (f (+ t (/ h 4))
                                       (add y
                                            (scalar-product 1/4 k1))))]
               [k3 (scalar-product h (f (+ t (* 3/8 h))
                                       (add y
                                            (scalar-product 3/32 k1)
                                            (scalar-product 9/32 k2))))]
               [k4 (scalar-product h (f (+ t (* 12/13 h))
                                       (add y
                                            (scalar-product 1932/2197 k1)
                                            (scalar-product -7200/2197 k2)
                                            (scalar-product 7296/2197 k3))))]
               [k5 (scalar-product h (f (+ t h)
                                       (add y
                                            (scalar-product 439/216 k1)
                                            (scalar-product -8 k2)
                                            (scalar-product 3680/513 k3)
                                            (scalar-product -845/4104 k4))))]
               [k6 (scalar-product h (f (+ t (/ h 2))
                                       (add y
                                            (scalar-product -8/27 k1)
                                            (scalar-product 2 k2)
                                            (scalar-product -3544/2565 k3)
                                            (scalar-product 1859/4104 k4)
                                            (scalar-product -11/40 k5))))]
               [R (/ (norm (add (scalar-product 1/360 k1)
                                (scalar-product -128/4275 k3)
                                (scalar-product -2197/75240 k4)
                                (scalar-product 1/50 k5)
                                (scalar-product 2/55 k6)))
                     h)]
               [delta (* 0.84 (expt (/ tol R) 1/4))])
          (if (<= R tol)
              (vector (+ t h)
                      (add y
                           (scalar-product 25/216 k1)
                           (scalar-product 1408/2565 k3)
                           (scalar-product 2197/4104 k4)
                           (scalar-product -1/5 k5))
                      (if (< delta 1) h
                          (* delta h)));step(>=h) as large as possible
              (if (< (* delta h) min-step)
                  (display "minimum h exceeded!");; error!!
                  (next-state (vector (vector-ref now-state 0)
                                      (vector-ref now-state 1)
                                      (* delta h))))))))
    (let iter ([now-state (vector init-time init-val init-step)]
               [result (list (vector init-time init-val init-step))])
      (if (end-condition now-state); end-condition is defined by user
          result
          (let ([new-state (next-state now-state)])
            (iter new-state (append result (list new-state))))))))

;;----test
;;choose an end-condition: time >= t, i.e. time in [t0,t]
(define end-condition
  (lambda (t)
    (lambda (state)
      (>= (vector-ref state 0) t))))
;;f = exp(2t) * sin(t) - 2*y0 + 2*y1 
(define (f t y)
  (list (list-ref y 1)
        (+ (* (exp (* 2 t)) (sin t))
           (* -2 (list-ref y 0))
           (* 2 (list-ref y 1)))))
;; precise solution: y0 = 0.2 * exp(2t) * (sin(t) - 2*cos(t))
(define (precise t) (* 0.2 (exp (* 2 t)) (- (sin t) (* 2 (cos t)))))
(define result-of-states (rkf f 0 (list -0.4 -0.6) 0.1 1e-2 (end-condition 8) 1e-3))
(define result-of-t-y0 (map (lambda (ele) (vector (vector-ref ele 0) (car (vector-ref ele 1)))) result-of-states))
;; compare precise and approximate solutions
(import (prefix gnuplot-pipe gp:))
(import list-comprehensions)
(define times (map (lambda (ele) (vector-ref ele 0)) result-of-t-y0))
(define vals (map precise times))
(define valsapprox (map (lambda (ele) (vector-ref ele 1)) result-of-t-y0))

(gp:call/gnuplot
 (gp:send "set xlabel 't'")
 ;; Plot multiple curves.
 (gp:plot (list (list "title 'precise'" times vals)
                (list "title 'numerical' with points" times valsapprox))
          ))
